[Frontend] Fix conv wrapper crash when the input is a view of a lower-rank buffer#294
Open
YWHyuk wants to merge 66 commits into
Open
[Frontend] Fix conv wrapper crash when the input is a view of a lower-rank buffer#294YWHyuk wants to merge 66 commits into
YWHyuk wants to merge 66 commits into
Conversation
… feed Skeleton + EmitC + cost/dep analysis on the frontend; the trace runtime, loader, bridge, and Core feed on the simulator; shared MLIR pass helpers and the pipeline tests.
Per-record tag key in the bridge plus per-iteration tag alloc in dma-fine-grained so multi-tile-K and conv loads do not collide; strip the reduction accum marker from the memory_barrier slot.
togsim_dispatch with TILE_BEGIN/TILE_END; outline each work-item into togsim_kernel_tile.
DMA-capacity throttle and frozen-state guard, per-core VMEM in the configs, and the SA weight-buffer throttle.
trace_timeline.py with per-work-item grouping and resource-centric DMA lanes; the trace logs the first DRAM response and the assigned systolic array, and scopes the compute barrier to its dispatch.
Default to the trace path; fix uninitialized Instruction fields, the matmul accumulator wedge, fused-subtile dedup, nested/fused epilogue dataflow, and dma_wait fusion; bound concurrent dispatches to the spad, round-robin work-items within a partition, benchmark autotune and run the multi-tenant scheduler through the trace path, and emit trace.so for pooling/reduction.
Carry simulator headers through the wrapper for cache-safe replay; drop verbose [P3-trace] logs; fix the key.mlir compile race in load().
… runtime model Replace the trace bridge's accumulated special cases with one dataflow rule and clean up the runtime that consumes it. Dependency rule: per SRAM buffer keep a writers SET; a reader depends on all current writers (occupancy=ISSUE when both are systolic-array ops, else latency=DONE); a writer REPLACEs the set. The only exception is is_mm_accum (a matmul that reads and writes the same buffer = a commutative accumulator): skip its read edge and UNION its write, waiting only the non-matmul init seed and not ordering co-matmuls. This drops the matmul-accumulator chain that deadlocked the SA weight-slot pipeline while keeping the init->matmul edge, and lets a vector epilogue or the store wait every K matmul (fixes the pure-vector store that an empty COMPUTE_BAR let slip). Remove COMPUTE_BAR entirely: a matmul is its own DONE-handle (finish == SA drain), so the store JOINs the matmul writers directly. The whole emit/loader chain is gone -- build_skeleton, lower_to_emitc, togsim.compute_barrier, the runtime symbol, the Opcode/case/_fence_finish, and TraceRec::COMPUTE_BAR -- so a stale producer fails loudly instead of emitting records the bridge would drop. Only MEMORY_BAR remains (an async load's DONE is its data arrival, not issue). Model compute-output spad footprint in the SRAM version/capacity machinery so buffer reuse (WAR) is capacity-modeled, not a hard edge. The output size comes from the DMA records that touch the same buffer (a buf_bytes pre-pass); an in-place buffer (accumulator, relu) is version-transparent so footprint is not double-counted. The occupy gate and version release sit in the MOVIN/MOVOUT/COMP issue points (release before the COMP skip path so a skipped matmul still frees). Runtime: collapse child_inst / _pipeline_children into one event-indexed _deps[ISSUE|DONE] with add_dep(c, on) and fire(e); collapse the weight-slot release queue and the async-load wakeup into one _due_events timed-effect table drained by process_due_events. Both are behavior-preserving (byte-identical). Require the weight-slot model: sa_weight_buffer_depth must be > 0 (errors at init), and the round-robin disable mode is removed. Degenerate traces (a consumer-less preload, an unpinned matmul) hit explicit error+exit guards rather than asserts that vanish under NDEBUG. Mark the legacy ONNX TOG path deprecated: it is superseded by the trace path, so TileGraphParser logs a deprecation warning and the TORCHSIM_LEGACY_TOG=1 opt-in warns at command build.
… spad/2
The validation-binary spad-overflow check sat inside `if functional_mode:`, so in
timing-only / autotune (non-functional) runs an over-spad tiling was never
rejected and reached TOGSim, which wedged ("spad too small") and crashed the
compile via assert 0. Move the compile + check out of the functional gate (the
Spike execution itself stays gated, run_spike below) and budget per dispatch at
spad/2 -- two work-items run concurrently (double buffer), so each must fit half
the spad or they deadlock competing for it. This matches the GEMM tiling gate
(max_spad_size = spad/2), which pointwise ops lacked. Fixes the resnet /
test_scheduler wedge where a fused BatchNorm+ReLU tile exceeded spad/2.
Every generated trace.cpp now opens with a "GENERATED ... DO NOT EDIT" banner carrying the TOGSim ABI version and the togsim_* call formats, so a dumped trace is self-documenting. Both are read from togsim_runtime.h at codegen time -- the version from the TOGSIM_ABI_VERSION define, the call-format text from a marked block next to the declarations -- so they never drift when the ABI changes.
…rint The bridge sums each work-item's distinct-buffer footprint (each buffer once, so a reduction's reloads of the same section do not inflate it) onto its Tile. can_issue then admits two concurrent dispatches only when each fits half the spad; a dispatch larger than spad/2 runs alone with the whole spad, so two work-items never compete for the shared spad and deadlock -- a runtime safety net beneath the codegen spad/2 gate. The footprint and resulting max_dispatch are logged on the TILE_SCHEDULED trace line for debugging.
…election select_tile fed the epilogue buffer count to gemm_combination_mapping as max(n_extra_read - 2, 0), which optimistically assumed the X and W DMA buffers were already freed and reusable by the epilogue. The codegen lays every buffer out as a disjoint .spad global and never reuses freed space, so the estimate undercounted the real footprint: for a fused matmul+relu kernel the budget came to 64512 B/lane (treated as a bare GEMM) while the emitted kernel used 89088 B/lane including the relu output buffer. Under the full-spad guard this was harmless (89088 < 131072), but the spad/2 guard rejected it and crashed test_transformer_fusion with SpadOverflowError, since the heuristic path has no tile-shrink fallback. Pass n_extra_node + n_extra_read instead: one output-tile-sized buffer per epilogue node plus one per extra read operand, matching what the codegen emits. For the matmul+relu kernel the budget now equals the actual footprint exactly, and tile selection picks TILE_M=128 (62976 B/lane) which fits spad/2. Liveness-based spad reuse and broadcast over-allocation are tracked separately as optimizations in issue #275.
The spad-overflow check summed the kernel stack frame into the per-lane scratchpad usage (spad_usage = stack_size + spad_size) and rejected the tiling when that exceeded the spad/2 budget. But only the .spad section actually lives in the scratchpad -- it is pinned there by the --section-start=.spad link option. The kernel stack is in main memory (sp is set up by pk in the -m region, not at the scratchpad vaddr), so it does not consume the scratchpad and must not be charged against it. The scalar frame is also shared across lanes, not per-lane, so adding it double-counted. On small configs (8x8) this falsely rejected feasible tilings: the wrapper3 conv2d/resnet/mistral kernels fit the 32 KB spad with room to spare but were tripped over the spad/2 gate purely by the ~200-800 B stack term, crashing compile with SpadOverflowError. Drop the stack term; the .spad-only check still correctly rejects real buffer overflows (e.g. sparsity, which is fixed separately by f05ac8a).
Spike v1.0.3 zero-inits the MVIN/MVOUT DMA address buffer so ROUNDUP padding entries are skipped instead of dereferencing uninitialized garbage, fixing the host store segfault on 8-lane (wrapper3) configs where a tile's split axis is not a multiple of vlane_stride*n_vu.
Add pytorchsim_functional_verify_per_kernel, a sub-option of pytorchsim_functional_mode that localizes the first compiled kernel whose Spike output diverges from a CPU reference. When enabled, the generated wrapper compares every realized buffer (the output of each fused kernel) against a CPU "golden" computed once per call by running the original aten graph (V.graph.module) on CPU with the same inputs, via an fx.Interpreter that records each node's output by name. A buffer is mapped to its originating fx node through V.graph.get_buffer().origin_node, so the check reports the kernel, the originating op, the offending indices and the max abs diff, then raises at the first divergence (stop-at-first). Comparison granularity is the fused-cluster output, the finest observable in a fused pipeline. Auto-disabled when functional mode is off (no Spike values to verify); the config accessor AND-gates the key with functional mode. Codegen bakes the verify_init/verify_check calls into the wrapper only when enabled at compile time, so clear the codegen cache when toggling. Tolerances via TORCHSIM_FUNCTIONAL_VERIFY_RTOL / _ATOL (default 1e-4). extension_functional_verify.py holds the graph registry and the runtime golden/compare logic; mlir_codegen_backend injects the calls at the wrapper level; extension_config reads the new YAML key.
The BMM tile selector fed only n_extra_node to gemm_combination_mapping and dropped the prologue's extra-read operands entirely, so a softmax-fused attention matmul (value^T @ softmax(scores)) was tiled as a bare BMM. The codegen lays the softmax max/sum operands out as their own disjoint weight-tile-sized .spad globals (buf3/buf4) and never reuses freed space, so the estimate undercounted the real footprint: on the 32x32 wrapper2 config (16 KB/lane spad/2 budget) tile selection picked TILE_N=512, whose emitted kernel used 16896 B/lane and overflowed the budget by 512 B, crashing test_transformer with SpadOverflowError. Under the 128x128 full-budget config the slack hid it. This mirrors the epilogue n_extra_read gap fixed for the GEMM template (commit f05ac8a), now on the BMM prologue path. Add an n_prologue_extra_read knob to gemm_combination_mapping that charges each extra prologue-read operand as one weight-tile-sized (TILE_K x TILE_N) .spad buffer, matching what the codegen emits, and have the BMM template count those operands the same way codegen does (the one numel-matching main input reuses the matmul-operand buffer; every other read gets its own global). Tile selection now rejects TILE_N=512 (16896 B/lane) and picks a fitting tile (8704 B/lane), so wrapper2 test_transformer passes the full Spike + allclose check. The new parameter defaults to 0, leaving the GEMM and conv callers unchanged.
Add docs/per-kernel-functional-verify.md (usage, options, mechanism, limitations, code map) and point to it from CLAUDE.md: a one-line entry in the debugging section and in the YAML knobs list.
Keep only a one-line note on what TOGSIM_ABI_VERSION guards; the per-bump v1..v12 history was noise nobody reads.
[Frontend] Budget fused-prologue spad buffers in BMM tile selection
…s off format_dma_inst_issued_trace_line and format_instruction_detail_line are only ever used as the message argument to trace_instruction_line (a spdlog::trace sink). Because the argument is evaluated eagerly at the call site, the kernel builds a formatted std::string for every issued/finished instruction and then spdlog drops it whenever the level is above trace (the default). Bail out of both formatters when trace logging is disabled so the fmt::format work is skipped entirely. This removes wasted work but is not the simulation-speed bottleneck (that is the per-cycle ready-queue rescan in Core::cycle, addressed separately); the conv kernel wall time is unchanged within noise.
Core::cycle() walked every tile's ready-instruction list on every simulated cycle to find one instruction to issue. Profiling an 8x8 conv kernel showed this scan is ~96% of simulation time and, within it, the walk is >99%: the ready list holds ~2000 instructions (mostly blocked on the SRAM-capacity / weight-slot throttle), only one issues per cycle, and the blocked ones are never removed, so the same ~2000 are re-examined every cycle (~2.6k list iterations/cycle, ~1e9 over 400k cycles). At small arrays this dominates because tiny tiles inflate both the ready-list length and the number of DMA-wait stall cycles. Gate the scan behind a per-Core _issue_dirty flag. The scan's outcome can only change when the ready set grows or a resource frees, so set the flag on exactly those events: - ready set grows: Instruction::dec_ready_counter, on enqueueing a now-ready instruction, sets the OWNING core's _issue_dirty via _owner_dirty_ref (wired in Core::issue when the tile is dispatched) -- per-core so a remote core's enqueue does not force every core to rescan; - SRAM freed (release_sram) and weight slot freed (apply_due) set _issue_dirty; - a new dispatch (issue) and a successful issue keep it set. On a cycle with none of these the scan would re-walk the same blocked instructions and issue nothing, so it is skipped. Instructions are never moved between queues, so there is no re-admission churn. Issue-identical: it never changes which instruction issues or when, only whether the (side-effect-free under those conditions) scan runs. The 8x8 conv kernel produces the same 403026 cycles as before, with its TOGSim wall dropping from 214.5s to 41.9s (~5x, 71% of scans skipped). A forced-scan invariant check (no skipped cycle ever issues or inline-finishes a zero-cycle COMP) found 0 violations across 39 kernels spanning GEMM, BMM, softmax and conv. DMA-bound kernels (the ones that hit the 6h CI cap) stall more and gain more.
In Core::cycle()'s issue scan, a COMP with compute_cycle == 0 finishes inline and calls instructions.erase(it), but does not set `issued` or break -- so control falls through to the `it++` at the end of the loop body, incrementing the std::list iterator that erase() just invalidated. That reads a freed node and is undefined behavior; it usually limps along because the freed node's next pointer is briefly intact, but under a different allocator / sanitizer / build it can crash or skip-or-duplicate the remaining ready instructions in that tile, corrupting issue order. Use the iterator erase() returns and continue, the standard erase-in-loop idiom. No behavior change on the path that happened to work; it just removes the UB.
…X TOG The main compile/sim path no longer generates or selects the legacy ONNX Tile-Operation-Graph. extension_codecache emits only trace.so + trace_cycles.tsv (the build-skip now keys on trace.so), and TOGSimulator.run_standalone always drives TOGSim with --trace_so. The TORCHSIM_LEGACY_TOG opt-in is removed from the frontend. The ONNX --models_list branch is kept solely for the STONNE sparse path (extension_op.py); TOGSim's C++ ONNX parser is untouched (separate PR). origins (which FX nodes a kernel came from) is preserved: logged per kernel run and recorded as a trailing "# origins:" line in trace_cycles.tsv -- the legacy ONNX TOG carried this as node metadata, and the C++ cycle-table loader stops at the comment so the current parser is unaffected. Also drop the dead tog_file param from mlir_gem5_compile_command, migrate scripts/chiplet.sh to --trace_so/--cycle_table (the trace path stubs per-tensor addresses and --attributes_list is no longer a Simulator option), and refresh the CLAUDE.md TOG-generation notes.
Rework per-kernel codegen so the loop body is an ordered list of load->compute->store steps (class Step + push_step + codegen_loops iteration) instead of the fixed dma_loads/compute/dma_stores buffers. A step may be compute-only or transfer-only; the compute_idx loop is emitted only for steps that actually have compute content. Drop the ad-hoc self.masks buffer: the reduction-tail mask (get_mask) now emits into the compute buffer, since a mask is just compute. Use the step model to express indirect (gather/scatter) access cleanly. When the indirect index is produced in the same pass, convert_indirect_indexing pushes a new step for the offset build, so the index load, the offset build and the gather become separate steps bridged through spad. This replaces the compute_dependecy -> dma_stores hack in both convert_indirect_indexing and load(). push_step clones the CSE so the name counter is shared but the dedup cache is per-step (each step is its own compute_idx region). Validated: x[idx+2]+1 emits two steps and matches CPU; add/matmul/reduce/softmax/layernorm and gather/scatter unchanged.
…ptor
Replace the affine.apply{indirect_access} symbol smuggle with an explicit
offset descriptor. convert_indirect_indexing returns the offset spad instead
of folding sympy.Symbol(out) into the index; emit_transfer carries it as a
togsim.transfer operand; decompose_transfer lifts that operand to a
memref.dma_start {indirect_offset = @spad_symbol} attribute (memref.dma_start
is a registered op and rejects an extra operand, but accepts the attribute);
lower_dma_to_gemmini reads the attribute and resolves the global for CONFIG4
(drops _find_indirect); build_skeleton adds the offset spad to the gather DMA
read_bufs so the offset-build -> gather dependency edge forms in the trace.
The index stays clean (base only). Validated on both paths: Spike functional
(computed-index gather + scatter allclose, pointwise/reduce regression 0) and
the C++ trace timing path end-to-end (allclose; gather togsim_dma read_bufs now
includes the offset spad).
Indirect addressing (scattered-DMA timing) in the new trace path is a separate
gap tracked in issue #284.
…ride, compute-loop multi-dim offset Three refinements on top of the explicit offset descriptor: 1. Detect indirect access by an explicit symbol set instead of an "tmp" substring match. indirect_indexing now records str(index_var) in self.indirect_symbols; a _has_indirect(expr) helper tests the index free_symbols against it; the former "tmp"-string sites (store, get_dma_info, convert entry/indirect_dims/stride) use it. Removes the now-dead _find_indirect from lower_dma_to_gemmini. 2. Single indirect dim: pass the raw index spad and let the MVIN apply the gather stride per position (CONFIG4 offset_stride) instead of a vector_load/muli/vector_store round-trip that pre-multiplied the stride. emit_transfer carries offset_stride; decompose copies the attribute; lower programs CONFIG4 with it (was hardcoded to 1). 3. Multi indirect dim (e.g. x[ix, iy]): the offset is a sum of strided indices, which a single CONFIG4 channel cannot do, so the sum stays in the kernel -- but build it in the compute loop (chunked by compute_vec_size, not a single tile-wide vector) and store it to a DEDICATED offset spad so an index that is live elsewhere (x[ix, iy] + ix) is not clobbered. push_step separates the offset-build loop from the gather that reads it. Adds multi-dim gather and index-reuse cases to test_indirect_access. Validated: indirect/scatter/embedding + the two new multi-dim cases pass; add/matmul/softmax regression 0.
…emmini memref.dma_start is a registered op with a fixed operand list, so it cannot carry the runtime descriptor operands the masked-DMA clamp needs (per-dim low/high are dynamic-shape values, not attributes). togsim.transfer is unregistered, so it carries them. Eliminate the dma_start intermediate: togsim.transfer flows through the whole pipeline and lowers directly to Gemmini. - New lower_transfer_to_gemmini merges decompose_transfer's <=4D handling (unit-dim collapse / >4D affine.for peel + lane-banked SRAM offset) with lower_dma_to_gemmini's CONFIG/CONFIG2/CONFIG3/[CONFIG4]/MVIN|MVOUT asm. A/B-validated (func7 + packed CONFIG rs1/rs2) against the old decompose->lower chain on add/pad/matmul/gather. - togsim.transfer gains a tag_idx operand (dram,dram_idx,sram,sram_idx,tag,tag_idx, dma_type,vst[,offset]) so the async DMA<->barrier runtime slot survives without memref.dma_start's variadic tag indices. - memref.dma_wait -> togsim.wait (togsim-level counterpart, emitted by the vcix matmul lowering, erased by the merged lowering). No memref.* DMA ops remain. - decompose_transfer dropped from PRE_OPT; loop-padding runs OPAQUELY on togsim.transfer (mlir-opt --allow-unregistered-dialect), doing only its loop-bound rounding (its DRAM-grow half is what the future masked-DMA clamp replaces). - dma_fine_grained, lower_to_vcix._DmaView, build_skeleton ported to read togsim.transfer. Functional path (Spike, timing_mode=0) validated: add, matmul (fine-grained subtile + async + togsim.wait), softmax, gather/scatter (indirect CONFIG4), constant_pad, conv (padding=0) all allclose. Known-pending: conv with padding>0 regresses -- it relied on loop-padding's DRAM-grow, which the masked-DMA clamp (next step) replaces (padding=0 passes; wrong values, not a crash).
…dims, cleanup Addresses the PR review of the masked-DMA non-dividing work: - _masked_fill_bits detected a log-sum-exp reduction with `"exp" in origin.name`, which also matches expand / expm1 / experimental -- an ordinary non-dividing sum whose origin name contains "exp" would then get a -inf tail fill instead of the sum identity 0 and reduce to -inf. Match the op target instead (_origin_is_exp: the target's overload packet name == "exp"); same fix applied to the get_padding_type twin. Unit test in tests/lowering/test_masked_fill.py pins that expand/expm1 do not match. - Remove the self._last_local_dims temporal coupling: get_dma_info now returns local_dims and load()/store() pass it to _masked_bounds explicitly, so a reordered or skipped get_dma_info can no longer feed stale clamp axes. - Drop the underscore-prefixed locals in def_dma_op / store_epilogue (non-idiomatic), trim the oversized _masked_fill_bits / _masked_bounds docstrings to the invariant plus a link to docs/design/masked-dma-descriptor.md, and clarify the bf16 dtype-table comment (present only for table totality; bf16 is not actually supported).
Dropping the tile-divisibility recompile means non-dividing shapes rely entirely on the masked-DMA clamp/fill; a gap is now a silent wrong answer, not a slow-but-safe recompile. Add tests/ops/misc/test_masked_nondividing.py pinning that path: 2-D constant pad (per-dim greedy pad recovery), sum over a broadcast operand (0 fill, guards the exp-substring bug), softmax (genuine -inf fill), non-dividing elementwise, and non-dividing gather. Register it in the CI allowlist (pytorchsim_test.yml). Also add the missing trailing newline to test_indirect_access.py.
…ilogue The ordered-steps refactor removed the reduction masking mechanism (self.masks / get_mask / mask_cse) but left a compute_body.splice(self.masks) in codegen_epilogue_body, so any matmul/bmm + reduction fusion crashed with AttributeError: no attribute 'masks'. It was masked until now by the stale-spike illegal-instruction CI failure. Remove the dead splice; test_matmul_reduction (matmul+max, var_mean, residual+var_mean) passes. Non-dividing reduction-fusion masking (padded lanes vs the reduction identity) stays a separate gap from that refactor.
…ixel_shuffle trace) Both the Gemmini descriptor and the TOGSim DMA model cap at 4 tile dims, but a logical tile can exceed 4D (e.g. pixel_shuffle splits two spatial axes -> a 5D tile [1,1,2,4,2]). lower_transfer_to_gemmini reduces >4D as part of emitting Gemmini asm, but that runs only on the Spike/gem5 lowering path. The trace producer (build_tog) reads togsim.transfer BEFORE that, so it emitted a 5D DMA and TOGSim rejected it: "issued tile is not supported format.. tile.size: 5". Exposed by the pixel_shuffle case newly added to the wrapper1 CI suite. Add a peel_transfer POST_OPT pass that reduces every >4D togsim.transfer up front (before build_tog) while KEEPING it a togsim.transfer, so both consumers see <=4D. It mirrors lower_transfer_to_gemmini's two reductions: collapse unit dims when the effective rank is <=4 (memref.collapse_shape), else peel the outer effective dims into an affine.for nest with the lane-banked physical SRAM offset. The gemmini lowering keeps its own reduction as a now-defensive no-op. Validated on 128x128: test_floormod_axis_split (pixel_shuffle + 9 others), test_masked_nondividing, test_matmul, test_conv2d all pass.
…-safe) _GoldenInterpreter recorded each node output on CPU but returned the original (npu) tensor to downstream nodes, and cast the recorded value to float32. A device-baked op (e.g. arange on npu) or a consumer that requires operand device agreement (aten.index needs its index tensors on the indexed tensor's device) then saw a mixed CPU/npu pair and the whole golden threw, silently disabling the per-kernel verify for that graph -- e.g. MobileNet's depthwise conv lowered to constant_pad + index gather. run_node now moves EVERY tensor output to CPU (preserving dtype) and returns it, so downstream nodes see CPU operands; only the recorded golden is cast to float32 for the allclose compare, so int index tensors stay usable.
The cat template's _calculate_input_tile_sizes handed each input a concat-dim tile of min(extent, remaining_spad_budget). When that budget did not cover an input's full concat extent, the tile width could fail to divide the extent (e.g. extent 64, tile 62 on the 8-lane config). The copy loop steps by that tile over [0, extent), so a non-dividing tile produces a ragged final iteration with no remainder mask: the MVIN over-reads past the input and the MVOUT over-writes past the output, corrupting the result. This surfaced on the 8x8 (wrapper3) CI config, where the per-input spad budget lands at 126 for two 64-wide inputs, truncating the second tile to 62. It scrambled ~45% of the cat output and broke LlamaDecoderLayer (rotary rotate_half cat) and the diffusion UNet skip-connection concat. Larger arrays give a bigger budget so both inputs got the full 64-wide (dividing) tile, which is why 32x32 and 128x128 passed. Snap each input's concat-dim tile down to the largest divisor of its extent that fits the budget, so every loop iteration is full-width and in-bounds. When the budget already covers the full extent the tile is unchanged, so 32x32/128x128 codegen is a no-op. Verified with a minimal cat repro (64+64 -> 128 on 8x8: fails before, diff 0 after) and the real test_llama LlamaDecoderLayer at 8x8 (now passes).
The vlane index_expr lowering reconstructs each lane element's logical index along the split axis as stride_dim + outer_dim * nr_vector_lane, where outer_dim is the sublane-row index. The next sublane row starts vector_lane * vlane_stride elements later (vector_lane lanes, each holding vlane_stride positions), but the multiplier used nr_vector_lane alone, dropping the vlane_stride factor. Any index tensor (e.g. arange/iota) whose split axis spills into a second sublane row then had every row past the first come out shifted, producing wrong values. This fired on the 8x8 (wrapper3) config: arange(32) at 8 lanes with vlane_stride 2 needs two sublane rows, so positions 16-31 came out as p-8. Larger arrays keep the 32-wide axis in a single row (outer_dim always 0), which is why 32x32/128x128 were unaffected. It broke test_llama's LlamaModel (rotary position ids) at 8x8. Multiply the outer/row term by vector_lane * vlane_stride. This equals the old value whenever vlane_stride == 1, so single-row splits and larger configs are a codegen no-op. Verified with a minimal arange(32) repro at 8x8 (positions 16-31 wrong before, diff 0 after) and the full test_llama LlamaModel at 8x8 (now passes).
init_tile_size gave a 1-D tile of 2*vlane_stride*vector_lane regardless of the data extent. When the extent is smaller than that tile the reduction/pointwise MVIN reads past the end of the DRAM buffer, and the compute loop folds the out-of-range elements into the result. For a sum reduction that means adjacent-DRAM garbage is squared/added into the sum. This broke test_single_perceptron's Loss on the 128x128 config: the MSE mean reduces 128 elements but the tile was 2*2*128 = 512, so the MVIN over-read 384 elements past the 128-wide input and roughly doubled the squared-error sum (loss diff ~4593). It fired on wide-lane configs (tile 512 > extent 128); 32x32 gives tile 128 == extent and 8x8 gives a tile that divides 128, so both were unaffected. Cap the 1-D tile to min(2*vlane_stride*vector_lane, extent) so the tile never exceeds the data. It is a no-op whenever extent >= tile, so larger reductions and other configs are unchanged. Verified with a minimal mse_loss(1-D) repro at 128x128 (over-reads before -> diff 0 after, across extents 64/100/127/200/300) and the real test_single_perceptron at 128x128 (Loss now passes).
torch.sort's frontend lowering exposes the sorted VALUES as the sole scheduler-visible template output; the INDICES buffer is written in place by the same kernel (make_inplace) but is not registered as an output. For argsort the values output is dead, so Inductor DCEs the whole sort kernel and the indices buffer is returned uninitialized (all-zeros), silently corrupting any argsort result (e.g. DeepSeek MoE expert routing). Register the in-place indices write as a MutationOutput owned by the sort operation so the scheduler keeps the kernel alive whenever indices is consumed. OperationBuffer.get_outputs() hardcodes [self], so add a small MLIRTemplateBuffer subclass that can own extra MutationOutputs and have MLIRTemplateCaller.output_node build it; the sort lowering attaches the indices mutation via add_mutation_output. For templates with no mutation this is behaviorally identical to the base (get_outputs == [self]). Then gate the values MVOUT on the values output being live: in argsort it is pruned from the kernel args, so emitting its MVOUT referenced an undeclared %YV. The values scratchpad is still produced internally for the bitonic compare; only the dead DRAM store is skipped. Verified: t.argsort() now returns correct indices (was all-zeros); torch.sort with both outputs live is unchanged; test_sort, test_matmul, test_cat, test_indirect_access unaffected by the output_node change. Note: x[argsort] end to end additionally needs the indirect-gather codegen fix (separate).
convert_indirect_indexing moves an indirect index into an offset spad and
zeros the indirect term in the DRAM base-offset expression, but it only
walked index.args. A BARE indirect index (x[idx], where the load index IS
the symbol so index.args is empty) was left unzeroed, so the DMA base
offset's affine.apply referenced the loop-local index SSA value before it
was defined ("use of undeclared SSA value name"). Any 1-D gather with a
runtime index tensor failed to compile at codegen; it was masked for
argsort only because the sort DCE bug left the indices all-zero (a
degenerate constant load).
Subs any remaining indirect symbol to 0 after the arg loop (mirrors the
store path). The per-position gather is carried by the offset spad, so the
base offset is correctly 0.
Verified: x[idx] (N=32/64) and x[x.argsort()] now compile and match CPU;
test_indirect_access.py (incl. scatter/index_add and multi-dim) and
test_sort.py unchanged.
_masked_bounds reverse-engineered per-dim padding from the single flat index offset and clamped each padded load to [pad, pad + input_extent). That recovery is ill-posed: the offset mixes the tap shift with the padding, and under channels_last the stride-1 channel axis absorbs the offset remainder, producing an impossible clamp (e.g. [4, 8) on a size-2 channel tile). Spike then zeroed the whole load, making channels_last depthwise conv 99.9 percent wrong while NCHW was fine. The clamp was also unnecessary: at pad positions the consumer already yields the correct value (a compute-side select for a padded gather, or the pad op's own fill), so reading OOB is harmless. Keep only the ragged tail clamp (glo=0, ghi=loop_extent), which is genuinely load-bearing for non-dividing tiles. Verified on 32x32: channels_last depthwise conv now matches CPU; NCHW conv and standalone ragged F.pad still pass. Broader regressions (padded reduction loads, full MobileNet) to be addressed separately.
codegen_epilogue_body() decided whether to emit the template buffer's MVOUT
by asking "did a fused epilogue already store something?":
if len(self.stores._lines) == 0:
template_store()
That is a proxy for the real question, "does anything outside this kernel still
read the template buffer?", and it is wrong in two of the four cases:
- epilogue + live template buffer -> store skipped, buffer never written
- no epilogue + dead template buffer -> store emitted for a pruned DRAM arg
The first case corrupts DeepSeek-V3's MoE gate. The gate is
sigmoid(mm(hidden, gate_w) + bias), so add+sigmoid fuse onto the GEMM template,
but the raw mm result is also read by a later gather kernel. Its buffer was
declared as a kernel output and never stored, so the gather read an
uninitialized buffer (all zeros), scrambling expert routing. Every config was
affected; only 8x8 crossed the test's loose atol=2e-1 (final Max abs diff
0.2546) while 32x32 landed just under it and passed.
Ask the real question instead. Inductor already answers it:
Kernel.remove_kernel_local_buffers() drops a buffer only when
Scheduler.can_buffer_be_removed_through_fusion() finds every user inside this
fused kernel. A buffer that survived removal therefore has an outside user and
must be stored. Record the template buffer's name in def_kernel() and
def_conv_kernel() (conv has its own implementation) and gate template_store()
on it, in both the main and the reduction_fusion branches. The proxy is gone.
Use the kernel-local self.removed_buffers, not V.graph.removed_buffers: the
former is what remove_buffer() populates, and the latter is still empty when
the store hook runs.
Note this makes a previously missing MVOUT issue, so DRAM traffic and cycle
counts rise for any model with a live template buffer under a fused epilogue.
That is the correct cost, not a regression, but reported cycles will shift.
Verified on the 8x8 config: test_deepseek_v3_base goes from Max abs diff 0.2546
to passing; a dead template buffer (relu(mm)) still prunes its DRAM arg and
emits a single MVOUT; the trace/TOGSim timing path runs with the extra store.
Regressions pass: test_bmm_reduction, test_matmul, test_bmm, test_conv2d,
test_cnn, test_group_conv, test_pool, test_reduce, test_softmax, test_layernorm,
test_mlp.
Implement abs, sign, isnan, isinf, fmod, bitwise shifts, copysign, erfc, hypot, cosh, sinh, acos, asin, atan, atan2, acosh, asinh, atanh, and log/ atan (VCIX) in the Inductor MLIR backend. Float-only transcendentals cast non-float inputs to f32 while leaving native float widths (f16/f64) untouched.
Port the C++ MathLog/MathAtanToVCIX patterns to the in-process Python VCIX pass: add math.log (opcode 2, imm 2) and math.atan (opcode 2, imm 3) to _MATH_VIV and MARKERS. Encoding matches the spike rs1 / gem5 VS1 sub-opcodes (sin=0, cos=1, log=2, atan=3). No C++ mlir-opt pass needed.
Add the pointwise op test under tests/ops/elementwise and register it in the test workflow allowlist.
Bump the third-party pins to the releases shipping the new pointwise instructions (spike torchsim_vlog/vatan, gem5 CustomVlog/Vatan + decode).
reduction_init returned "-inf"/"inf" for max/min regardless of dtype, so an
integer max/min reduction fed an invalid inf identity into every consumer:
the accumulator init, the template reduction init, and the masked-DMA tail
fill (torch.tensor(inf, dtype=int64) overflowed -- surfaced compiling swinv2
with batch>1). reduction_combine_vec also hardcoded the float multi_reduction
kinds <maximumf>/<minimumf>, which reject integer vectors.
Return the dtype's representable extreme from reduction_init for integer
max/min (max identity -inf -> iinfo.min, min identity +inf -> iinfo.max;
bool -> 0/1), and emit the signed-integer multi_reduction kinds <maxsi>/<minsi>
for integer element types. sum/prod need no change: the <add>/<mul> combining
kinds are polymorphic and ops.add/ops.mul already select arith.add{i,f} by
dtype (as ops.maximum/minimum already select maxsi vs maximumf).
Verified: int amax/amin/sum/prod reductions compile and match CPU; float
test_reduce (ReduceMax) and test_softmax unchanged. Also clears the swinv2
batch>1 masked-fill overflow (it then hits a separate unsupported
shifted-window view, tracked separately).
…lar DMA index
torch.roll has no Inductor lowering; it decomposes to index_select(fmod(...)) -- a
cyclic-shift gather whose ModularIndexing DMA index the affine-only descriptor
cannot express ("Unlinearized floor/mod in DMA index"). Rewrite it as the two
contiguous halves the shift produces, stitched by cat:
roll(x, s, d) == cat([slice(x, N-s, N), slice(x, 0, N-s)], d)
each half is a plain strided slice with no modulo. Do it as a lowering (removing
roll from the decomp table so it is not decomposed first) rather than a
decomposition, so we can also REALIZE the input with copy_input: roll's producer
is often a reshaped view (e.g. swinv2's window_reverse), and the slice offset would
otherwise fuse into that reshape's modular index, re-creating a shifted (p+c)%m
form. A plain .contiguous()/clone does NOT create the buffer boundary (Inductor
inlines it); copy_input does.
Verified: torch.roll (1-D / multi-dim / arbitrary shift) and roll+reduction match
CPU; swinv2 batch>1 now compiles through the shifted-window attention (previously
blocked by the roll's cyclic shift and then by the fused slice offset).
A composition of aligned reshapes on one iteration variable (e.g. swinv2's window partition fused with a later reshape) leaves a nested index like ModularIndexing(ModularIndexing(p, 1, 64), 1, 8) that neither sympy nor simplify_with_ranges reduces. collect_boundaries then skips its cut points (the inner base is not a bare variable) and the affine-only DMA check rejects it. Add a general, pattern-free canonicalizer (flatten_nested_floormod / _as_digit): any nesting of FloorDiv/ModularIndexing over a single symbol is a digit extractor (v // A) % M and collapses to one level by composing divisors, from four divisibility-guarded algebraic identities. Applied in collect_boundaries and in _fold_with_ranges. Multi-variable inner arguments (e.g. a roll shift v+c) are left untouched. Verified: numeric equivalence on the swinv2 window index; test_reduce, test_layernorm, test_cat, test_indirect_access unaffected.
The wrapper header imported empty_strided but not empty_strided_cpu, so a graph that allocates a CPU-side buffer (e.g. swinv2's shifted-window attention mask, which Inductor keeps on CPU) failed at runtime with "NameError: name 'empty_strided_cpu' is not defined". Add the same binding the default Inductor PythonWrapperCodegen header provides.
…yout A per-lane reduction is lowered as a 2-D [reduction | batch] collapse: multi_reduction reshapes the accumulator to <reduction_numel x red_size> and reduces axis 0, which assumes the reduction axis is the outermost in-lane run and every non-reduction (batch) axis is inner. get_dma_info reorders the tile axes for reduction loads to guarantee that, but only for the 2-D and 3-D cases. The 4-D case mis-tested is_reduction (reduction_depth < 3, false once there are three batch dims) and then raised NotImplementedError, and rank 5+ fell through to the default row-major order -- both leaving the reduction axis innermost, so the reduce collapsed a batch axis instead of the reduction axis. This is hit whenever a non-reduction dim-merge is blocked and an extra batch axis stays in-lane: SwinV2 cosine window attention adds a gathered relative-position bias (table[idx[q,k], head]) before the softmax amax, keeping head separate from query (4-D tile, head bled into the max), and its post-attention LayerNorm var_mean splits the token axis into several loop vars (6-D tile, mean off by ~0.005). Generalize the reorder to any tile rank: a tile that carries the reduction axis places that axis-group outermost via range(r, L) + range(r-1, -1, -1), which reduces to the existing 3-D order for L=3. Adds test_reduce_gather_bias (fails max abs diff ~3.3 before the fix).
SwinV2 batch>1 used to fail codegen with "Unlinearized floor/mod in DMA index" because the shifted-window path (torch.roll composed with window partition/reverse) produced views axis-split could not linearize. With the roll->slice+cat lowering, the nested/shifted mod handling, and the reduction-axis layout fix, the whole backbone now compiles and matches CPU end to end (max diff ~5e-6 for image 64 / window 8 / batch 2). Add tests/models/test_swinv2.py, wire it as a self-hosted CI job next to the other model tests, and list SwinV2 in the README model coverage.
conv_multi_tile_mapping (used for batch > 1 convs) collects every tile that fits SPAD into tile_candidates and returns them sorted, but it guarded the "no mapping" error on max_used_spad_size instead of on tile_candidates. max_used_spad_size is only bumped when a candidate also satisfies max_k_h_w <= k_h, and max_k_h_w is initialized to K_W, so it only ever fires for the full-kernel (k_h == K_H) tile. When that tile overflows SPAD -- e.g. a CLIP ViT-B/32 patch conv (3->768, 32x32 stride 32) at batch > 1 -- the guard raised "Cannot find a valid mapping" even though smaller-k_h tiles fit and were already in tile_candidates. The mapping variable it tracks is never returned. Guard on `not tile_candidates` instead, so the conv is rejected only when no tile fits at all. The returned candidate list is unchanged, so convs that already worked are unaffected. Adds a batched patch-conv case to test_conv2d.py. Fixes issue #252.
With the conv-mapping fix, CLIP's vision transformer runs with batch > 1 and matches CPU end to end (max diff ~1e-5). Add tests/models/test_clip.py (CLIPVisionModel, batch 2, patch 32), wire it as a self-hosted CI job, and list CLIP in the README model coverage.
A conv input node can be a ReinterpretView. call_kernel passes template args
by buffer *name*, so the wrapper receives the base buffer, whose rank/shape may
differ from the view codegen assumed. The wrapper then did
padded_shape = list(X.shape)
padded_shape[3] += 2 * PADDING_W # IndexError on a 3D base buffer
This fires whenever the view is free: a contiguous (B, N, C) buffer already is
the channels_last layout of (B, C, H, W), so inductor inserts no materializing
copy. That is exactly the transformer (B, N, C) -> (B, C, H, W) pattern, e.g.
SegFormer efficient attention.
Stop inferring geometry from the runtime tensor. Bake the size/stride/offset the
codegen used into the wrapper and rebuild the logical NCHW inputs from it, so the
wrapper is correct whether it is handed the base buffer, an already reinterpreted
view, or a materialized copy (reinterpret is idempotent in all three cases). The
conv templates were the only place reading .shape at runtime; everything else
already derives geometry from the node layout at codegen time.
No change to call_kernel, mlir_argdefs, the MLIR signature, or arg_attributes.
Also allocate the padding buffer directly on the input device with the input
dtype: torch.zeros() defaulted to f32 and to(device=...) round-tripped via CPU.
Covers a conv fed by a free ReinterpretView of a lower-rank buffer. Fails with IndexError on the pre-fix wrapper and passes after it.
d9131ba to
1ff2ebb
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes a conv codegen bug where the generated wrapper crashes with
Root cause
A conv input node can be a
ReinterpretView.call_kernelpasses template args bybuffer name, so the wrapper is handed the base buffer, whose rank/shape may differ
from the view codegen assumed:
extract_info)X.layout.size = [1, 32, 8, 8](the 4D view)X.shape = (1, 64, 32)(the base buffer)The wrapper then inferred geometry from the runtime tensor (
list(X.shape),X.shape[2],X.shape[3]) and blew up.It only fires when the view is free: a contiguous
(B, N, C)buffer already is thechannels_last layout of
(B, C, H, W), so inductor inserts no materializing copy. That isexactly the transformer
(B, N, C) -> (B, C, H, W)pattern — e.g. SegFormer'sefficient-attention spatial-reduction conv (
Conv2d(C, C, k=sr, stride=sr), padding 0).When the layout does not match, inductor inserts a copy kernel, a real 4D buffer arrives,
and the bug is invisible. That is why it went unnoticed.
Fix
Stop inferring geometry from the runtime tensor. Bake the size/stride/offset the codegen
used into the wrapper and rebuild the logical NCHW inputs from it:
X.layoutdescribes the data within the storage of the tensor that arrives (both come fromthe same node), so the reinterpret is correct — and idempotent — whether the wrapper is
handed the base buffer, an already reinterpreted view, or a materialized copy.
The conv templates were the only place reading
.shapeat runtime; everything else alreadyderives geometry from the node layout at codegen time. No change to
call_kernel,mlir_argdefs, the MLIR signature, orarg_attributes— the call site still passes theraw base buffer and it now works.
Also fixed for free: the padding buffer was allocated via
torch.zeros(shape)(always f32)then
.to(device=...)(CPU round-trip). It now allocates directly on device withX.dtype.Verification
tests/ops/conv/test_conv_view_input.py: fails with the aboveIndexErroron the pre-fix wrapper, passes after (max diff 2.3e-05 / 2e-04 vs CPU).padding=1) unchanged and numerically correct(max diff 3.8e-05 vs CPU);
reinterpret_tensoris an identity there.Follow-ups (not in this PR)
padding == 0fast path: shapes are literals now, so thezerosalloc + full copy can beskipped for 1x1 / strided convs.
mlir_argdefsderivesbuffer_typesnumel from the base buffer, so a slicing viewfeeding e.g. gemm can still get a wrong MLIR arg size. Same class of bug, separate fix.
def_conv_kernelpatches the signature withre.sub(r'(\d+)(?=xf32)', ...), hardcodingxf32; an f16 conv would not get its padded input size applied.Found while checking #254 on this branch: SegFormer no longer hits the originally reported
Not supporting formaterror, but was blocked here instead.